#!/usr/bin/env python3
import cv2
import subprocess
import signal
import sys
import threading
import time
import numpy as np

class Cam2RTMP:
    def __init__(self,
                 src=0,
                 rtmp_url='rtmp://8.155.164.202/live/livestream?secret=17aa39b9cbc2453a82261c8afdae9436',
                 width=640,
                 height=480,
                 fps=30):
        self.src = src
        self.rtmp_url = rtmp_url
        self.w = width
        self.h = height
        self.fps = fps

        # 打开摄像头
        self.cap = cv2.VideoCapture(src)
        if not self.cap.isOpened():
            raise RuntimeError(f'无法打开摄像头 {src}')
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
        self.cap.set(cv2.CAP_PROP_FPS, fps)

        # 最新帧与锁
        self.latest_frame = None
        self.lock = threading.Lock()

        # 启动 FFmpeg 子进程
        self.ffmpeg = None
        self._start_ffmpeg()

        # 线程标志位
        self.running = True

        # 采集线程
        self.capture_thread = threading.Thread(target=self._grab_loop, daemon=True)
        self.capture_thread.start()

        print('🚀 摄像头推流已启动，按 Ctrl-C 停止')

    # -------------------------- #
    def _start_ffmpeg(self):
        cmd = [
            'ffmpeg',
            '-y',
            '-f', 'rawvideo',
            '-vcodec', 'rawvideo',
            '-pix_fmt', 'bgr24',
            '-s', f'{self.w}x{self.h}',
            '-r', str(self.fps),
            '-i', '-',
            '-c:v', 'libx264',
            '-pix_fmt', 'yuv420p',
            '-preset', 'ultrafast',
            '-tune', 'zerolatency',
            '-f', 'flv',
            self.rtmp_url
        ]
        self.ffmpeg = subprocess.Popen(cmd, stdin=subprocess.PIPE)
        print(f'🎥 推流地址: {self.rtmp_url}')

    # -------------------------- #
    def _grab_loop(self):
        """独立线程循环读摄像头"""
        while self.running:
            ret, frame = self.cap.read()
            if not ret:
                print('❌ 摄像头读取失败')
                continue
            # 分辨率统一
            if (frame.shape[1], frame.shape[0]) != (self.w, self.h):
                frame = cv2.resize(frame, (self.w, self.h))
            with self.lock:
                self.latest_frame = frame

    # -------------------------- #
    def spin(self):
        """主线程：按 FPS 把帧写入 FFmpeg"""
        period = 1.0 / self.fps
        while self.running:
            with self.lock:
                frame = self.latest_frame
            if frame is None:
                time.sleep(0.001)
                continue
            try:
                self.ffmpeg.stdin.write(frame.tobytes())
            except BrokenPipeError:
                print('❌ FFmpeg 断开，停止推流')
                self.running = False
                break
            time.sleep(period)

    # -------------------------- #
    def stop(self):
        print('🛑 正在停止……')
        self.running = False
        self.capture_thread.join()
        if self.ffmpeg:
            try:
                self.ffmpeg.stdin.close()
                self.ffmpeg.terminate()
                self.ffmpeg.wait()
            except Exception:
                pass
        self.cap.release()
        print('👋 已清理完毕')

# ================================================================ #
def main():
    src = int(sys.argv[1]) if len(sys.argv) > 1 else 0
    url = sys.argv[2] if len(sys.argv) > 2 else 'rtmp://8.155.164.202/live/livestream?secret=17aa39b9cbc2453a82261c8afdae9436'

    streamer = Cam2RTMP(src=src, rtmp_url=url)

    def _sigint_handler(sig, frame):
        streamer.stop()
        sys.exit(0)

    signal.signal(signal.SIGINT, _sigint_handler)

    try:
        streamer.spin()
    except KeyboardInterrupt:
        pass
    finally:
        streamer.stop()

if __name__ == '__main__':
    main()